clear
load ../parameter_sandbox_wage_rule/eq_mit_mass.mat 

v2    = reshape(eq_SS.v.v2, glob.n(1), glob.n(2)); % assume that the expected value is in SS at date T

tic() 

for t = options.T:-1:1
   param.M = sim.M(t);
    
    param.C = sim.C(t);
    param.D = sim.D(t);
    param.w = sim.W(t); % HH FOC
    param.A = sim.A(t); 
    
    param.mu_f = sim.cF(t);
    param.ce = sim.cE(t);
    
    c = reshape(v2, glob.n(1), glob.n(2));
    c = [c; c];

    
    % First solve the incumbent's value function
    v     = solve_valfunc(c,glob.s,param,glob,options);
    
    %% Solve entrant problem on the grid of signals
    
    c2           = c(glob.Nb+1:end,:);
    c2           = reshape(c2, glob.n(1), glob.n(2));
    v2           = griddedInterpolant(glob.B, glob.S,c2, 'linear');   
    
    %% Compute entry choice and entrant policy
                   
    % compute entrant policy
    B                   = [glob.minb*ones(glob.n(2),1), glob.maxb*ones(glob.n(2),1)];
    obj                 = @(y)valfunc_E(v2, glob.agrid,y,param,glob,options);     
    y                   = goldenx(obj,B(:,1),B(:,2));

    [~,ve]              = valfunc_E(v2, glob.agrid,y,param,glob,options);
    eq.ve   = ve.v1;
    
    entry_value   = griddedInterpolant(glob.agrid, eq.ve, 'linear');
    entry_value   = entry_value(glob.sig_grid);

    prob_entry    = normcdf(entry_value - param.ce, 0, param.ce_noise);
    sim.entry(:,t) = prob_entry;

    sim.l_e(:,t)        = ve.l;

    %% continue with incumbent problem
    v2    = reshape(v.v2, glob.n(1), glob.n(2));

    % Solve on a finer grid to get policy functions
    tmp      = glob.PnK;
    glob.PnK = glob.PnKf;
    vf       = solve_valfunc(c,glob.sf,param,glob,options);
    glob.PnK = tmp;
 
    % Record optimal policy
    sim.y(:,t)  = vf.y; 
    sim.l(:,t)  = vf.l;
    sim.p(:,t)  = vf.p;  
    
    sim.exit(:,t) = vf.p_exit;
    
    if (t == 1)
        disp('hi')
    end
    
end
toc()

%% Now iterate forwards on policies
tic()

fspaceerg   = fundef({'spli',glob.bgridf,0,1});


for t = 1:options.T
    
    param.M = sim.M(t);

    % Distribution from last period
   if t > 1
       L               = sim.mu(:,t-1);
       
       lp = sim.l(:,t-1);
       le = sim.l_e(:,t-1);
       
       p_enter         = sim.entry(:,t);
       p_exit          = sim.exit(:,t-1);
   else
        L               = eq_SS.L;
        
        lp = eq_SS.l;
        le = eq_SS.ve_l;
        
        p_enter         = eq_SS.prob_entry;
        p_exit          = eq_SS.p_exit;
   end
    

   
    % set up G matrix
    entrant_dist  = glob.Qshift'*(p_enter.*glob.a_stat_E');

    G                  = zeros(glob.Nsf,1);
    G(find(glob.sf(:,1) == glob.bgridf(end))) = entrant_dist;
    
    entrant_policy = repmat(le, 1, glob.nf(1))';
    entrant_policy = entrant_policy(:)';
    
    Qb             = funbas(fspaceerg, entrant_policy');

    Qent           = dprod(glob.QeI,Qb);
    G = Qent'*G;
   
    % set up transition matrix
    lp          = min(lp, glob.maxb);
    lp          = max(lp, glob.minb);
      
    Qb              = funbas(fspaceerg, lp);
    Q               = dprod(glob.Qef,Qb);
        
    sim.mu(:,t) = Q'*((1-param.gamma)*(1-p_exit).*L)+ param.M*Q'*G; 

    sim.entry_mass(t) = sum(param.M*G);
    sim.entry_rate(t) = sim.entry_mass(t)/sum(sim.mu(:,t));

    eq_mit.L          = sim.mu(:,t);
    eq_mit.yf         = sim.y(:,t);
    
    aggs              = find_agg_mit(eq_mit, param, glob, options);
    
    sim.C(t)          = aggs.C;
    sim.D(t)          = aggs.D;

    L = sum(eq_mit.L.*lp);
    
    
    % distribution of firms ages a=0,1,2,3,4,5
    sim.a0(:,t)      = param.M*Q'*G;
    
    if t == 1
        sim.a1(:,t)      = eq_SS.Q'*((1-param.gamma)*(1-p_exit).*sim.a0(:,t));
        sim.a2(:,t)      = eq_SS.Q'*((1-param.gamma)*(1-p_exit).*sim.a1(:,t));
        sim.a3(:,t)      = eq_SS.Q'*((1-param.gamma)*(1-p_exit).*sim.a2(:,t));
        sim.a4(:,t)      = eq_SS.Q'*((1-param.gamma)*(1-p_exit).*sim.a3(:,t));
        sim.a5(:,t)      = eq_SS.Q'*((1-param.gamma)*(1-p_exit).*sim.a4(:,t));
    else  
        sim.a1(:,t)      = Q'*((1-param.gamma)*(1-sim.exit(:,t-1)).*sim.a0(:, t-1));
        sim.a2(:,t)      = Q'*((1-param.gamma)*(1-sim.exit(:,t-1)).*sim.a1(:, t-1));
        sim.a3(:,t)      = Q'*((1-param.gamma)*(1-sim.exit(:,t-1)).*sim.a2(:, t-1));
        sim.a4(:,t)      = Q'*((1-param.gamma)*(1-sim.exit(:,t-1)).*sim.a3(:, t-1));
        sim.a5(:,t)      = Q'*((1-param.gamma)*(1-sim.exit(:,t-1)).*sim.a4(:, t-1));
    end
    
    if t == 6
        disp('hello, it is t = 6')
    end
    
    sim.employment_a0(t) = sum(lp.*sim.a0(:,t));
    sim.employment_a1(t) = sum(lp.*sim.a1(:,t));
    sim.employment_a2(t) = sum(lp.*sim.a2(:,t));
    sim.employment_a3(t) = sum(lp.*sim.a3(:,t));
    sim.employment_a4(t) = sum(lp.*sim.a4(:,t));
    sim.employment_a5(t) = sum(lp.*sim.a5(:,t));

    sim.employment_young(t) = sim.employment_a0(t) + sim.employment_a1(t) + ...
                              sim.employment_a2(t) + sim.employment_a3(t) + ...
                              sim.employment_a4(t);
    
                              % employment at firms of age < 1
    sim.entrant_employment(t) = sim.employment_a0(t);

                          
    rev = sim.p(:,t).*sim.y(:,t);
                          
    sim.output_a0(t) = sum(rev.*sim.a0(:,t));
    sim.output_a1(t) = sum(rev.*sim.a1(:,t));
    sim.output_a2(t) = sum(rev.*sim.a2(:,t));
    sim.output_a3(t) = sum(rev.*sim.a3(:,t));
    sim.output_a4(t) = sum(rev.*sim.a4(:,t));
    sim.output_a5(t) = sum(rev.*sim.a5(:,t));
    
    sim.output_young(t) = sim.output_a0(t) + sim.output_a1(t) + ...
                              sim.output_a2(t) + sim.output_a3(t) + ...
                              sim.output_a4(t) + sim.output_a5(t);
    
    sim.ngdp(t) = sum(rev.*sim.mu(:,t));
                          
    sim.L(t) = L;

    
    sim.W(t)          = L^param.nu*param.psi*(sim.C(t));
    
end

employment_young = sim.employment_young;
save('employment_young.mat', 'employment_young')
